﻿#region Png Error Checking Recommendations
/*
10.1. Error checking
To ensure early detection of common file-transfer problems, decoders should verify that all eight bytes of the PNG file signature are correct. (See Rationale: PNG file signature.) A decoder can have additional confidence in the file's integrity if the next eight bytes are an IHDR chunk header with the correct chunk length. 
Unknown chunk types must be handled as described in Chunk naming conventions. An unknown chunk type is not to be treated as an error unless it is a critical chunk. 
It is strongly recommended that decoders should verify the CRC on each chunk. 
In some situations it is desirable to check chunk headers (length and type code) before reading the chunk data and CRC. The chunk type can be checked for plausibility by seeing whether all four bytes are ASCII letters (codes 65-90 and 97-122); note that this need be done only for unrecognized type codes. If the total file size is known (from file system information, HTTP protocol, etc), the chunk length can be checked for plausibility as well. 
If CRCs are not checked, dropped/added data bytes or an erroneous chunk length can cause the decoder to get out of step and misinterpret subsequent data as a chunk header. Verifying that the chunk type contains letters is an inexpensive way of providing early error detection in this situation. 
For known-length chunks such as IHDR, decoders should treat an unexpected chunk length as an error. Future extensions to this specification will not add new fields to existing chunks; instead, new chunk types will be added to carry new information. 
Unexpected values in fields of known chunks (for example, an unexpected compression method in the IHDR chunk) must be checked for and treated as errors. However, it is recommended that unexpected field values be treated as fatal errors only in critical chunks. An unexpected value in an ancillary chunk can be handled by ignoring the whole chunk as though it were an unknown chunk type. (This recommendation assumes that the chunk's CRC has been verified. In decoders that do not check CRCs, it is safer to treat any unexpected value as indicating a corrupted file.) 
*/
#endregion

#region Imported Libraries
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using org.webguvenligi.NSecureImageLib.GraphicFileFormats.Utilities;
#endregion

namespace org.webguvenligi.NSecureImageLib.GraphicFileFormats.PNG
{
    public class PNGParser : BaseImageParser
    {

        public PNGParser(System.IO.Stream bitmapStream)
            : base(bitmapStream)
        {
        }

        private static byte[] PNGHeader =
            new byte[] { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };

        public List<PNGBaseChunk> ResultChunks
        {
            get;
            private set;
        }

        private int ChunkTypeCount(Type type)
        {
            if (!typeof(PNGBaseChunk).IsAssignableFrom(type))
            {
                throw new ArgumentException("type", "Invalid Type passed");
            }
            int chunkCount = (from chunk in this.ResultChunks
                              where chunk.GetType() == type
                              select chunk).Count();
            return chunkCount;
        }
        private readonly string[] notMultipleAllowedChunks =
             new string[] { "IHDR", "PLTE", "IEND", "cHRM", "gAMA", "iCCP", "sBIT", "sRGB", "bKGD", "hIST", "tRNS", "pHYs", "tIME", "sCAL" };
        /// <summary>
        /// Inserts the specified target chunk.
        /// </summary>
        /// <param name="targetChunk">The target chunk.</param>
        public void Insert(PNGBaseChunk targetChunk)
        {
            if (targetChunk == null)
                throw new ArgumentNullException("targetChunk");
            if (this.ResultChunks.FindChunks(targetChunk.GetType()).FirstOrDefault() == null)
                throw new ArgumentException("Target Chunk Type can occur only once");
            //Ancillary Chunks Should be located before IDAT Chunks
            if (targetChunk is PNGAncillaryChunkBase)
            {
                int chunkIndex = 0;
                this.ResultChunks.FindChunk<IDATChunk>(out chunkIndex);
                ResultChunks.Insert(chunkIndex, targetChunk);
            }
        }
        public override void WriteImage(Stream targetStream,
                               bool removeMetadataTags)
        {
            bool concatIDATChunks = false;
            if (targetStream != null &&
                targetStream.CanWrite &&
                targetStream.CanSeek)
            {
                targetStream.Position = 0;
                using (BinaryWriter targetStreamWriter = new BinaryWriter(targetStream))
                {
                    targetStreamWriter.Write(PNGHeader);

                    List<PNGBaseChunk>.Enumerator chunkEnum =
                        ResultChunks.GetEnumerator();

                    List<IDATChunk> idatChunks =
                         new List<IDATChunk>();

                    while (chunkEnum.MoveNext())
                    {
                        PNGBaseChunk tempChunk =
                            chunkEnum.Current;

                        if (concatIDATChunks &&
                            tempChunk.GetType() == typeof(IDATChunk))
                        {
                            ConcatenateIDATChunks(concatIDATChunks,
                                                  ref chunkEnum,
                                                  idatChunks,
                                                  ref tempChunk,
                                                  targetStreamWriter);
                        }
                        if (!(removeMetadataTags &&
                              typeof(PNGAncillaryChunkBase).IsAssignableFrom(tempChunk.GetType())))
                        {
                            tempChunk.WriteChunkData(targetStreamWriter);
                        }
                    }
                    targetStreamWriter.Flush();
                }
            }
        }

        private void ConcatenateIDATChunks(bool concatIDATChunks,
                                           ref List<PNGBaseChunk>.Enumerator chunkEnum,
                                           List<IDATChunk> idatChunks,
                                           ref PNGBaseChunk tempChunk,
                                           BinaryWriter chunkWriter)
        {
            if (concatIDATChunks)
            {
                double totalChunkSize = 0;

                int chunkCount =
                    this.ChunkTypeCount(typeof(IDATChunk));

                if (tempChunk.GetType() == typeof(IDATChunk))
                {
                    totalChunkSize +=
                        tempChunk.TotalChunkSize;
                    idatChunks.Add((IDATChunk)tempChunk);
                    while (chunkEnum.MoveNext()
                           && chunkEnum.Current.GetType() == typeof(IDATChunk))
                    {
                        tempChunk = chunkEnum.Current;
                        totalChunkSize +=
                             tempChunk.TotalChunkSize;
                        idatChunks.Add((IDATChunk)tempChunk);
                    }
                    tempChunk = chunkEnum.Current;
                }

                int optimumChunkSize =
                    CalculateOptimumChunkSize(totalChunkSize);

                if (chunkCount > optimumChunkSize)
                {
                    OptimizeChunks(idatChunks, chunkWriter, totalChunkSize, optimumChunkSize);
                }
                else
                {
                    for (int i = 0; i < idatChunks.Count; i++)
                    {
                        idatChunks[i].WriteChunkData(chunkWriter);
                    }
                }
            }
        }

        private static void OptimizeChunks(List<IDATChunk> idatChunks,
                                           BinaryWriter chunkWriter,
                                           double totalChunkSize,
                                           int optimumChunkSize)
        {
            if (optimumChunkSize == 1)
            {
                Crc32Hash tempHash = new Crc32Hash();
                chunkWriter.WriteLittleEndian((uint)totalChunkSize);
                chunkWriter.Write(IDATHeaderConstant);
                tempHash.Update(IDATHeaderConstant);
                for (int i = 0; i < idatChunks.Count; i++)
                {
                    byte[] bodyData = idatChunks[i].InternalChunkData;
                    tempHash.Update(bodyData);
                    chunkWriter.Write(bodyData);
                }
                chunkWriter.Write(tempHash.CrcByteArray);
                chunkWriter.Flush();
            }
            else
            {
                List<IDATChunk>.Enumerator chunkReader =
                    idatChunks.GetEnumerator();
                Crc32Hash tempHash = new Crc32Hash();
                int currentChunkLength = 0;
                int remainder = 0;
                for (int i = 0; i < optimumChunkSize; i++)
                {
                    if (remainder == 0)
                    {
                        chunkWriter.WriteLittleEndian(PNGBaseChunk.MaxChunkSize);
                        chunkWriter.Write(IDATHeaderConstant);
                        tempHash.Update(IDATHeaderConstant);
                    }
                    byte[] sourceBytes = null;
                    int bytesToRead = 0;
                    while (currentChunkLength
                            <= PNGBaseChunk.MaxChunkSize)
                    {
                        chunkReader.MoveNext();
                        IDATChunk currentChunk =
                            chunkReader.Current;
                        bytesToRead =
                            (int)currentChunk.TotalChunkSize;
                        if ((currentChunkLength + currentChunk.TotalChunkSize) >
                            PNGBaseChunk.MaxChunkSize)
                        {
                            bytesToRead =
                                PNGBaseChunk.MaxChunkSize - currentChunkLength;
                            remainder =
                                (int)currentChunk.TotalChunkSize - bytesToRead;
                        }
                        currentChunkLength += bytesToRead;
                        sourceBytes = currentChunk.InternalChunkData;
                        tempHash.Update(sourceBytes, 0, bytesToRead);
                        chunkWriter.Write(sourceBytes, 0, bytesToRead);
                    }
                    currentChunkLength = 0;
                    chunkWriter.Write(tempHash.CrcByteArray);
                    tempHash.Reset();
                    if (remainder > 0)
                    {
                        chunkWriter.WriteLittleEndian(PNGBaseChunk.MaxChunkSize);
                        chunkWriter.Write(IDATHeaderConstant);
                        tempHash.Update(IDATHeaderConstant);
                        chunkWriter.Write(sourceBytes, bytesToRead, remainder);
                        tempHash.Update(sourceBytes, bytesToRead, remainder);
                        currentChunkLength = remainder;
                    }
                }
            }
        }

        private static readonly byte[] IDATHeaderConstant =
                            Encoding.ASCII.GetBytes("IDAT");

        private static int CalculateOptimumChunkSize(double totalChunkSize)
        {
            int optimumChunkSize =
                Convert.ToInt32((totalChunkSize / Convert.ToDouble(PNGBaseChunk.MaxChunkSize)));
            optimumChunkSize +=
                (totalChunkSize % (double)PNGBaseChunk.MaxChunkSize > 0 ? 1 : 0);
            return optimumChunkSize;
        }

        /// <summary>
        /// Parses the image.
        /// </summary>
        /// <returns></returns>
        protected override bool ParseImage()
        {
            using (BinaryReader pngReader = new BinaryReader(ImageStream))
            {
                ResultChunks = PNGBaseChunk.GetChunks(pngReader).ToList();
            }
            return true;
        }

        /// <summary>
        /// Cleans the additional chunks.
        /// </summary>
        /// <returns></returns>
        public int CleanAdditionalChunks()
        {
            return ResultChunks.RemoveAll(new Predicate<PNGBaseChunk>(delegate(PNGBaseChunk test)
            {
                bool isCorrectType = typeof(PNGAncillaryChunkBase).IsAssignableFrom(test.GetType());
                if (isCorrectType)
                    System.Diagnostics.Debug.WriteLine(String.Format("{0} Chunk Removed", test.ChunkTypeValue.ChunkTypeDescription));
                return isCorrectType;
            }));
        }

        /// <summary>
        /// Gets the format exception message.
        /// </summary>
        /// <value>The format exception message.</value>
        protected override string FormatExceptionMessage
        {
            get { return "Not a valid png file"; }
        }

        protected override byte[] HeaderSignature
        {
            get { return PNGHeader; }
        }
    }
}